import os
import os.path as osp
import xml.etree.ElementTree as ET

import mmcv
import numpy as np
import cv2
from DOTA_devkit.ssdd_evaluation import voc_eval
# from mmdet.core import norm_angle
# from mmdet.core import rotated_box_to_poly_single,poly_to_rotated_box_single
from .builder import DATASETS
from .xml_style import XMLDataset
from mmdet.core import reval_map, rdets2points
def rotated_box_to_poly_single(rrect):
    """
    rrect:[x_ctr,y_ctr,w,h,angle]
    to
    poly:[x0,y0,x1,y1,x2,y2,x3,y3]
    """
    x_ctr, y_ctr, width, height, angle = rrect[:5]
    tl_x, tl_y, br_x, br_y = -width / 2, -height / 2, width / 2, height / 2
    rect = np.array([[tl_x, br_x, br_x, tl_x], [tl_y, tl_y, br_y, br_y]])
    R = np.array([[np.cos(angle), -np.sin(angle)],
                  [np.sin(angle), np.cos(angle)]])
    poly = R.dot(rect)
    x0, x1, x2, x3 = poly[0, :4] + x_ctr
    y0, y1, y2, y3 = poly[1, :4] + y_ctr
    poly = np.array([x0, y0, x1, y1, x2, y2, x3, y3], dtype=np.float32)
    poly = get_best_begin_point_single(poly)
    return poly
def poly_to_rotated_box_single(poly):
    """
    poly:[x0,y0,x1,y1,x2,y2,x3,y3]
    to
    rotated_box:[x_ctr,y_ctr,w,h,angle]
    """
    poly = np.array(poly[:8], dtype=np.float32)

    pt1 = (poly[0], poly[1])
    pt2 = (poly[2], poly[3])
    pt3 = (poly[4], poly[5])
    pt4 = (poly[6], poly[7])

    edge1 = np.sqrt((pt1[0] - pt2[0]) * (pt1[0] - pt2[0]) +
                    (pt1[1] - pt2[1]) * (pt1[1] - pt2[1]))
    edge2 = np.sqrt((pt2[0] - pt3[0]) * (pt2[0] - pt3[0]) +
                    (pt2[1] - pt3[1]) * (pt2[1] - pt3[1]))

    width = max(edge1, edge2)
    height = min(edge1, edge2)

    angle = 0
    if edge1 > edge2:
        angle = np.arctan2(
            np.float(pt2[1] - pt1[1]), np.float(pt2[0] - pt1[0]))
    elif edge2 >= edge1:
        angle = np.arctan2(
            np.float(pt4[1] - pt1[1]), np.float(pt4[0] - pt1[0]))
    print(angle)
#     angle = norm_angle(angle)

    x_ctr = np.float(pt1[0] + pt3[0]) / 2
    y_ctr = np.float(pt1[1] + pt3[1]) / 2
    rotated_box = np.array([x_ctr, y_ctr, width, height, angle])
    return rotated_box

@DATASETS.register_module()
class SSDDDataset(XMLDataset):
    CLASSES = ('ship',)

    def load_annotations(self, ann_file):
        """Load annotation from XML style ann_file.

        Args:
            ann_file (str): Path of XML file.

        Returns:
            list[dict]: Annotation info from XML file.
        """

        data_infos = []
        img_ids = mmcv.list_from_file(ann_file)
        self.img_names = []
        for img_id in img_ids:
            filename = f'AllImages/{img_id}.jpg'
            xml_path = osp.join(self.img_prefix, 'Annotations',
                                f'{img_id}.xml')
            tree = ET.parse(xml_path)
            root = tree.getroot()
            size = root.find('size')
            width = int(size.find('width').text)
            height = int(size.find('height').text)
            self.img_names.append(img_id)
            data_infos.append(
                dict(id=img_id, filename=filename, width=width, height=height))

            
            
            
            
        return data_infos

#     def get_ann_info(self, idx):
#         """Get annotation from XML file by index.

#         Args:
#             idx (int): Index of data.

#         Returns:
#             dict: Annotation info of specified index.
#         """

#         img_id = self.data_infos[idx]['id']
#         xml_path = osp.join(self.img_prefix, 'Annotations', f'{img_id}.xml')
#         data_infos = []
        
# #         if not xml_path:  # test phase, use image folder as ann_folder to generate pseudo annotations
# #             ann_files = glob.glob(ann_folder + '/*.png')
# #             for ann_file in ann_files:
# #                 data_info = dict()
# #                 img_id = osp.split(ann_file)[1][:-4]
# #                 img_name = img_id + '.png'
# #                 data_info['filename'] = img_name
# #                 data_info['ann'] = dict()
# #                 data_info['ann']['bboxes'] = []
# #                 data_info['ann']['labels'] = []
# #                 data_infos.append(data_info)
                
#         tree = ET.parse(xml_path)
#         root = tree.getroot()
#         gt_bboxes = []
#         gt_labels = []
#         gt_polygons = []
#         bboxes_ignore = []
#         labels_ignore = []
        
#         data_info = dict()
#         data_info['ann_info'] = dict()
#         for obj in root.findall('object'):
#             name = obj.find('name').text
#             label = self.cat2label[name]
#             difficult = int(obj.find('difficult').text)
#             bnd_box = obj.find('polygon')
#             poly = [
#                 int(bnd_box.find('x1').text),
#                 int(bnd_box.find('y1').text),
#                 int(bnd_box.find('x2').text),
#                 int(bnd_box.find('y2').text),
#                 int(bnd_box.find('x3').text),
#                 int(bnd_box.find('y3').text),
#                 int(bnd_box.find('x4').text),
#                 int(bnd_box.find('y4').text),
#             ]
            
#             print('-----------------------------')

#             rbbox = cv2.minAreaRect(np.array(poly, dtype=np.float32).reshape([4, 2]))
#             x, y, w, h, a = rbbox[0][0], rbbox[0][1], rbbox[1][0], rbbox[1][1], rbbox[2]
#             if w == 0 or h == 0:
#                 continue
#             while not 0 > a >= -90:
#                 if a >= 0:
#                     a -= 90
#                     w, h = h, w
#                 else:
#                     a += 90
#                     w, h = h, w
#             a = a / 180 * np.pi
#             assert 0 > a >= -np.pi / 2


#             gt_bboxes.append([x, y, w, h, a])
#             gt_labels.append(label)
#             gt_polygons.append(poly)
            
#         if gt_bboxes:
#         data_info['ann_info']['bboxes'] = np.array(
#             gt_bboxes, dtype=np.float32)
#         data_info['ann_info']['labels'] = np.array(
#             gt_labels, dtype=np.int64)
#         data_info['ann_info']['polygons'] = np.array(
#             gt_polygons, dtype=np.float32)
# #         else:
# #             data_info['ann_info']['bboxes'] = np.zeros(
# #                     (0, 5), dtype=np.float32)
# #             data_info['ann_info']['labels'] = np.array([], dtype=np.int64)
# #             data_info['ann_info']['polygons'] = np.zeros(
# #                         (0, 8), dtype=np.float32)


#         data_info['ann_info']['bboxes_ignore'] = np.zeros(
#                 (0, 5), dtype=np.float32)
#         data_info['ann_info']['labels_ignore'] = np.array(
#                 [], dtype=np.int64)
#         data_info['ann_info']['polygons_ignore'] = np.zeros(
#                 (0, 8), dtype=np.float32)

#         data_infos.append(data_info)
            
            
            
            
# #-------------------------------------------------            
# #             bbox = [x, y, w, h, a]
            
# #             ignore = False
# #             if self.min_size:
# #                 assert not self.test_mode
# #                 if bbox[2] < self.min_size or bbox[3] < self.min_size:
# #                     ignore = True
# #             if difficult or ignore:
# #                 bboxes_ignore.append(bbox)
# #                 labels_ignore.append(label)
# #             else:
# #                 bboxes.append(bbox)
# #                 labels.append(label)
                
                
# #         if not bboxes:
# #             bboxes = np.zeros((0, 5))
# #             labels = np.zeros((0,))
# #         else:
# #             bboxes = np.array(bboxes, ndmin=2)
# #             labels = np.array(labels)
# #         if not bboxes_ignore:
# #             bboxes_ignore = np.zeros((0, 5))
# #             labels_ignore = np.zeros((0,))
# #         else:
# #             bboxes_ignore = np.array(bboxes_ignore, ndmin=2)
# #             labels_ignore = np.array(labels_ignore)
# #         ann = dict(
# #             bboxes=bboxes.astype(np.float32),
# #             labels=labels.astype(np.int64),
# #             bboxes_ignore=bboxes_ignore.astype(np.float32),
# #             labels_ignore=labels_ignore.astype(np.int64))
#         return data_infos
    def get_ann_info(self, idx):
        """Get annotation from XML file by index.

        Args:
            idx (int): Index of data.

        Returns:
            dict: Annotation info of specified index.
        """

        img_id = self.data_infos[idx]['id']
        xml_path = osp.join(self.img_prefix, 'Annotations', f'{img_id}.xml')
        tree = ET.parse(xml_path)
        root = tree.getroot()
        bboxes = []
        labels = []
        polys = []
        bboxes_ignore = []
        labels_ignore = []
        polys_ignore = []
        for obj in root.findall('object'):
            name = obj.find('name').text
            label = self.cat2label[name]
            difficult = int(obj.find('difficult').text)
            bnd_box = obj.find('polygon')
            poly = [
                int(bnd_box.find('x1').text),
                int(bnd_box.find('y1').text),
                int(bnd_box.find('x2').text),
                int(bnd_box.find('y2').text),
                int(bnd_box.find('x3').text),
                int(bnd_box.find('y3').text),
                int(bnd_box.find('x4').text),
                int(bnd_box.find('y4').text),
            ]
        
#             bbox = poly_to_rotated_box_single(poly)
#----------------------------------------------------------------------------------------            
            rbbox = cv2.minAreaRect(np.array(poly, dtype=np.float32).reshape([4, 2]))
            x, y, w, h, a = rbbox[0][0], rbbox[0][1], rbbox[1][0], rbbox[1][1], rbbox[2]
            if w == 0 or h == 0:
                continue
            while not 0 > a >= -90:
                if a >= 0:
                    a -= 90
                    w, h = h, w
                else:
                    a += 90
                    w, h = h, w
            a = a / 180 * np.pi
            assert 0 > a >= -np.pi / 2
            bbox = [x, y, w, h, a]
            
#             print(a)
#----------------------------------------------------------------------------------------            
            ignore = False
            if self.min_size:
                assert not self.test_mode
                if bbox[2] < self.min_size or bbox[3] < self.min_size:
                    ignore = True
            if difficult or ignore:
                bboxes_ignore.append(bbox)
                labels_ignore.append(label)
                polys_ignore.append(poly)
            else:
                bboxes.append(bbox)
                labels.append(label)
                polys.append(poly)
        if not bboxes:
            bboxes = np.zeros((0, 5))
            labels = np.zeros((0,))
            polys = np.zeros((0, 8))
        else:
            bboxes = np.array(bboxes, ndmin=2)
            labels = np.array(labels)
            polys = np.array(polys)
        if not bboxes_ignore:
            bboxes_ignore = np.zeros((0, 5))
            labels_ignore = np.zeros((0,))
            polys_ignore = np.zeros((0, 8))
        else:
            bboxes_ignore = np.array(bboxes_ignore, ndmin=2)
            labels_ignore = np.array(labels_ignore)
            polys_ignore = np.array(polys_ignore)
        ann = dict(
            bboxes=bboxes.astype(np.float32),
            labels=labels.astype(np.int64),
            polygons = polys.astype(np.float32),
            bboxes_ignore=bboxes_ignore.astype(np.float32),
            labels_ignore=labels_ignore.astype(np.int64),
            polygons_ignore = polys_ignore.astype(np.float32))

        return ann
    def get_cat_ids(self, idx):
        """Get category ids in XML file by index.

        Args:
            idx (int): Index of data.

        Returns:
            list[int]: All categories in the image of specified index.
        """

        cat_ids = []
        img_id = self.data_infos[idx]['id']
        xml_path = osp.join(self.img_prefix, 'Annotations', f'{img_id}.xml')
        tree = ET.parse(xml_path)
        root = tree.getroot()
        for obj in root.findall('object'):
            name = obj.find('name').text
            label = self.cat2label[name]
            cat_ids.append(label)

        return cat_ids
    
    def evaluate(self,
                     results,
                     metric='mAP',
                     logger=None,
                     proposal_nums=(100, 300, 1000),
                     iou_thr=0.5,
                     scale_ranges=None):
            """Evaluate the dataset.

            Args:
                results (list): Testing results of the dataset.
                metric (str | list[str]): Metrics to be evaluated.
                logger (logging.Logger | None | str): Logger used for printing
                    related information during evaluation. Default: None.
                proposal_nums (Sequence[int]): Proposal number used for evaluating
                    recalls, such as recall@100, recall@1000.
                    Default: (100, 300, 1000).
                iou_thr (float | list[float]): IoU threshold. It must be a float
                    when evaluating mAP, and can be a list when evaluating recall.
                    Default: 0.5.
                scale_ranges (list[tuple] | None): Scale ranges for evaluating mAP.
                    Default: None.
            """
            if not isinstance(metric, str):
                assert len(metric) == 1
                metric = metric[0]
            allowed_metrics = ['mAP']
            if metric not in allowed_metrics:
                raise KeyError(f'metric {metric} is not supported')
            annotations = [self.get_ann_info(i) for i in range(len(self))]
            eval_results = {}
            if metric == 'mAP':
                assert isinstance(iou_thr, float)
                mean_ap, _ = reval_map(
                    results,
                    annotations,
                    scale_ranges=scale_ranges,
                    iou_thr=iou_thr,
                    dataset=self.CLASSES,
                    logger=logger)
                eval_results['mAP'] = mean_ap
            return eval_results

    def _det2str(self, results):
            mcls_results = {cls: '' for cls in self.CLASSES}
            for idx in range(len(self)):
                img_id = self.img_ids[idx]
                result = results[idx]
                for label in range(len(result)):
                    bboxes = rdets2points(result[label])
                    cls_name = self.CLASSES[label]
                    for i in range(bboxes.shape[0]):
                        resstr = '{:s} {:.12f} {:.1f} {:.1f} {:.1f} {:.1f} {:.1f} {:.1f} {:.1f} {:.1f}\n'
                        ps = list(bboxes[i][:-1])
                        score = float(bboxes[i][-1])
                        resstr = resstr.format(img_id, score, *ps)
                        mcls_results[cls_name] += resstr
            return mcls_results

    def _results2submission(self, results, out_folder=None):
            dota_results = self._det2str(results)
            if out_folder is not None:
                os.makedirs(out_folder, exist_ok=True)
                for cls in dota_results:
                    fname = f'Task1_{cls}.txt'
                    fname = os.path.join(out_folder, fname)
                    with open(fname, 'w') as f:
                        f.write(dota_results[cls])
            return dota_results

    def format_results(self, results, submission_dir=None, **kwargs):
            """Format the results to submission text (standard format for DOTA evaluation).

            Args:
                results (list): Testing results of the dataset.
                submission_dir (str | None): The folder that contains submission files. 
                    If not specified, a temp folder will be created. Default: None.

            Returns:
                tuple: (result_files, tmp_dir), result_files is a dict containing
                    the json filepaths, tmp_dir is the temporal directory created
                    for saving json files when submission_dir is not specified.
            """
            assert isinstance(results, list), 'results must be a list'
            assert len(results) == len(self), (
                'The length of results is not equal to the dataset len: {} != {}'.
                format(len(results), len(self)))

            if submission_dir is None:
                submission_dir = tempfile.TemporaryDirectory()
            else:
                tmp_dir = None
            result_files = self._results2submission(results, submission_dir)
            return result_files, tmp_dir
#     def evaluate(self, results, work_dir=None, gt_dir=None, imagesetfile=None,logger=None,):
#         results_path = osp.join(work_dir, 'results_txt')
#         mmcv.mkdir_or_exist(results_path)

#         print('Saving results to {}'.format(results_path))
#         self.result_to_txt(results, results_path)

#         detpath = osp.join(results_path, '{:s}.txt')
#         annopath = osp.join(gt_dir, '{:s}.xml')  # data/HRSC2016/Test/Annotations/{:s}.xml

#         classaps = []
#         map = 0
#         for classname in self.CLASSES:
#             rec, prec, ap = voc_eval(detpath,
#                                      annopath,
#                                      imagesetfile,
#                                      classname,
#                                      ovthresh=0.5,
#                                      use_07_metric=False)
# #------------------计算F1,Prec,Recall-------------------------------------------------            
#             F1 = 2*rec*prec/(prec+rec+1e-5)
#             rounded_prec = [ '%.4f' % elem for elem in prec ]
#             rounded_rec = [ '%.4f' % elem for elem in rec ]
#             rounded_F1 = [ '%.4f' % elem if elem<100 else 0 for elem in F1 ]
#             max_index = rounded_F1.index(max(rounded_F1))
#             max_F1 = rounded_F1[max_index]
#             max_prec = rounded_prec[max_index]
#             max_rec = rounded_rec[max_index]
         
#             print("Prec: ",max_prec)
#             print("Recall:",max_rec)
#             print("F1: ",max_F1)   
# #------------------计算F1,Prec,Recall-------------------------------------------------             
#             map = map + ap
#             print(classname, ': ', ap)
#             classaps.append(ap)

#         map = map / len(self.CLASSES)
#         print('map:', map)
#         classaps = 100 * np.array(classaps)
#         print('classaps: ', classaps)
#         # Saving results to disk
#         with open(osp.join(work_dir, 'eval_results.txt'), 'w') as f:
#             res_str = 'mAP:' + str(map) + '\n'
#             res_str += 'classaps: ' + ' '.join([str(x) for x in classaps])
#             f.write(res_str)
#         return map

#     def result_to_txt(self, results, results_path):
#         img_names = [data_info['id'] for data_info in self.data_infos]

#         assert len(results) == len(img_names), 'len(results) != len(img_names)'

#         for classname in self.CLASSES:
#             f_out = open(osp.join(results_path, classname + '.txt'), 'w')
#             print(classname + '.txt')
#             # per result represent one image
#             for img_id, result in enumerate(results):
#                 for class_id, bboxes in enumerate(result):
#                     if self.CLASSES[class_id] != classname:
#                         continue
#                     if bboxes.size != 0:
#                         for bbox in bboxes:
#                             score = bbox[5]
#                             bbox = rotated_box_to_poly_single(bbox[:5])
#                             temp_txt = '{} {:.4f} {:.4f} {:.4f} {:.4f} {:.4f} {:.4f} {:.4f} {:.4f} {:.4f}\n'.format(
#                                 osp.splitext(img_names[img_id])[0], score, bbox[0], bbox[1], bbox[2], bbox[3], bbox[4],
#                                 bbox[5], bbox[6], bbox[7])
#                             f_out.write(temp_txt)
#             f_out.close()
